// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/memory/arbitration_operation.h>
#include <mutex>

#include <pollux/common/base/exceptions.h>
#include <pollux/common/base/runtime_metrics.h>
#include <pollux/common/memory/memory.h>
#include <pollux/common/testutil/test_value.h>
#include <pollux/common/time/timer.h>

using kumo::pollux::common::testutil::TestValue;

namespace kumo::pollux::memory {
    using namespace kumo::pollux::memory;

    ArbitrationOperation::ArbitrationOperation(
        ScopedArbitrationParticipant &&participant,
        uint64_t requestBytes,
        uint64_t timeoutNs)
        : requestBytes_(requestBytes),
          timeoutNs_(timeoutNs),
          createTimeNs_(getCurrentTimeNano()),
          participant_(std::move(participant)) {
        POLLUX_CHECK_GT(requestBytes_, 0);
    }

    ArbitrationOperation::~ArbitrationOperation() {
        POLLUX_CHECK_NE(
            state_,
            State::kRunning,
            "Unexpected arbitration operation state on destruction");
    }

    std::string ArbitrationOperation::stateName(State state) {
        switch (state) {
            case State::kInit:
                return "init";
            case State::kWaiting:
                return "waiting";
            case State::kRunning:
                return "running";
            case State::kFinished:
                return "finished";
            default:
                return fmt::format("unknown state: {}", static_cast<int>(state));
        }
    }

    void ArbitrationOperation::setState(State state) {
        switch (state) {
            case State::kWaiting:
                POLLUX_CHECK_EQ(state_, State::kInit);
                break;
            case State::kRunning:
                POLLUX_CHECK(this->state_ == State::kWaiting || state_ == State::kInit);
                break;
            case State::kFinished:
                POLLUX_CHECK_EQ(this->state_, State::kRunning);
                break;
            default:
                POLLUX_UNREACHABLE(
                    "Unexpected state transition from {} to {}", state_, state);
                break;
        }
        state_ = state;
    }

    void ArbitrationOperation::start() {
        POLLUX_CHECK_EQ(state_, State::kInit);
        participant_->startArbitration(this);
        setState(ArbitrationOperation::State::kRunning);
        POLLUX_CHECK_EQ(startTimeNs_, 0);
        startTimeNs_ = getCurrentTimeNano();
    }

    void ArbitrationOperation::finish() {
        setState(State::kFinished);
        POLLUX_CHECK_EQ(finishTimeNs_, 0);
        finishTimeNs_ = getCurrentTimeNano();
        participant_->finishArbitration(this);
    }

    bool ArbitrationOperation::aborted() const {
        return participant_->aborted();
    }

    uint64_t ArbitrationOperation::executionTimeNs() const {
        if (state_ == State::kFinished) {
            POLLUX_CHECK_GE(finishTimeNs_, createTimeNs_);
            return finishTimeNs_ - createTimeNs_;
        } else {
            const auto currentTimeNs = getCurrentTimeNano();
            POLLUX_CHECK_GE(currentTimeNs, createTimeNs_);
            return currentTimeNs - createTimeNs_;
        }
    }

    bool ArbitrationOperation::hasTimeout() const {
        return state_ != State::kFinished && timeoutNs() <= 0;
    }

    uint64_t ArbitrationOperation::timeoutNs() const {
        if (state_ == State::kFinished) {
            return 0;
        }
        const auto execTimeNs = executionTimeNs();
        if (execTimeNs >= timeoutNs_) {
            return 0;
        }
        return timeoutNs_ - execTimeNs;
    }

    void ArbitrationOperation::setGrowTargets() {
        // We shall only set grow targets once after start execution.
        POLLUX_CHECK_EQ(state_, State::kRunning);
        POLLUX_CHECK(
            maxGrowBytes_ == 0 && minGrowBytes_ == 0,
            "Arbitration operation grow targets have already been set: {}/{}",
            succinctBytes(maxGrowBytes_),
            succinctBytes(minGrowBytes_));
        participant_->getGrowTargets(requestBytes_, maxGrowBytes_, minGrowBytes_);
        POLLUX_CHECK_LE(requestBytes_, maxGrowBytes_);
    }

    ArbitrationOperation::Stats ArbitrationOperation::stats() const {
        POLLUX_CHECK_EQ(state_, State::kFinished);
        POLLUX_CHECK_NE(startTimeNs_, 0);

        const uint64_t executionTimeNs = this->executionTimeNs();

        POLLUX_CHECK_GE(startTimeNs_, createTimeNs_);
        const uint64_t localArbitrationWaitTimeNs = startTimeNs_ - createTimeNs_;
        if (globalArbitrationStartTimeNs_ == 0) {
            return {
                localArbitrationWaitTimeNs,
                finishTimeNs_ - startTimeNs_,
                0,
                executionTimeNs
            };
        }

        POLLUX_CHECK_GE(globalArbitrationStartTimeNs_, startTimeNs_);
        const uint64_t localArbitrationExecTimeNs =
                globalArbitrationStartTimeNs_ - startTimeNs_;
        return {
            localArbitrationWaitTimeNs,
            localArbitrationExecTimeNs,
            finishTimeNs_ - globalArbitrationStartTimeNs_,
            executionTimeNs
        };
    }

    std::ostream &operator<<(std::ostream &out, ArbitrationOperation::State state) {
        out << ArbitrationOperation::stateName(state);
        return out;
    }
} // namespace kumo::pollux::memory
